#!/usr/bin/env python
# -*- coding: UTF-8 -*-
#
# Copyright [Gtlions Lai].
# Create Date:
# Update Date:
__authors__ = '"Gtlions Lai" <gtlions.l@qq.com>'

"""Summary about the module or script.


Summary about function.


  Class: about class.
  Function(): about function.
"""

import time
import sys
import platform

if platform.system() == 'Windows' or platform.system() == 'Darwin':
    from psycopg2 import *
elif platform.system() == 'Linux':
    from pygresql.pgdb import *
import datetime
import dsn


def partition_add_day_date(relation, reserve):
    relation = str(relation)
    reserve = int(reserve)
    db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
    cur = db.cursor()
    cur.execute("select to_char(current_date,'yyyy-mm-dd')")
    partitionrangestart = cur.fetchone()[0]
    cur.execute("select to_char(current_date+1,'yyyy-mm-dd')")
    partitionrangeend = cur.fetchone()[0]
    cur.execute("select count(*) from pg_tables where schemaname||'.'||tablename=lower('" + relation + "')")
    v_cnt = cur.fetchone()[0]
    if v_cnt:
        pass
    else:
        print 'ERROR: relation "' + relation + '" does not exist'
        return
    if reserve > 0:
        pass
    else:
        print 'ERROR: parameter "i_partition_add_cnt" must greater than 0,but you input ' + str(reserve)
        return

    cur.execute(
        "select pk_name_d,pk_name_d_suffix,p_d_parameter from gpmg.manager_table where schemaname||'.'||tablename=lower('" + relation + "')")
    ret = cur.fetchone()
    pk_name_d = ret[0]
    pk_name_d_suffix = ret[1]
    p_d_parameter = ret[2]
    sql_exists_partitions = "select 'alter table " + relation + " add partition '|| partitionname ||' start ('||partitionrangestart||') end ('||partitionrangeend||') " + p_d_parameter + ";' from pg_partitions where partitionschemaname||'.'||tablename='" + relation + "' and to_date(substr(partitionrangestart,2,11),'yyyy-mm-dd')>='" + partitionrangestart + "'::date"
    cur.execute("select to_char(current_date,'" + pk_name_d_suffix + "')")
    p_name_d = pk_name_d + cur.fetchone()[0]
    sql_next_partitions = ""
    for i in range(1, reserve + 1):
        sql_next_partitions += "select 'alter table " + relation + " add partition " + p_name_d + " start (''" + partitionrangestart + " 00:00:00''::timestamp without time zone) end (''" + partitionrangeend + " 00:00:00''::timestamp without time zone) " + p_d_parameter + ";'"
        if i < reserve:
            sql_next_partitions += " union all "
        cur.execute("select to_char(current_date+" + str(i) + ",'" + pk_name_d_suffix + "')")
        pk_name_d_suffix_1 = cur.fetchone()[0]
        p_name_d = pk_name_d + pk_name_d_suffix_1
        cur.execute("select to_char(current_date+" + str(i) + ",'yyyy-mm-dd')")
        partitionrangestart = cur.fetchone()[0]
        cur.execute("select to_char(current_date+ " + str(i + 1) + ",'yyyy-mm-dd')")
        partitionrangeend = cur.fetchone()[0]
    cur.execute("select nextval('gpmg.seq_manager_part_mid')")
    manager_part_mid_sid = cur.fetchone()[0]
    cur.execute(
        "insert into gpmg.manager_part_mid(sid,t1) select " + str(
            manager_part_mid_sid) + ",* from (" + sql_exists_partitions + ") a ;")
    cur.execute(
        "insert into gpmg.manager_part_mid(sid,t2) select " + str(
            manager_part_mid_sid) + ",* from (" + sql_next_partitions + ") a ;")
    cur.execute(
        "insert into gpmg.manager_part_mid(sid,t3) select sid,t2 from gpmg.manager_part_mid where sid=" + str(
            manager_part_mid_sid) + " except select sid,t1 from gpmg.manager_part_mid where sid=" + str(
            manager_part_mid_sid) + ";")
    cur.execute("select t3 from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid) + " and t3 is not null")
    v_msg = "Done.Nothing."
    for record in cur.fetchall():
        sql_add_partitions = record[0]
        cur.execute(sql_add_partitions)
        v_msg = v_msg.replace('Nothing.', '')
    cur.execute("delete from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid))
    cur.close()
    db.commit()
    db.close()
    return v_msg


def partition_split_day_date(relation, reserve):
    relation = str(relation)
    reserve = int(reserve)
    db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
    cur = db.cursor()
    cur.execute("select to_char(current_date,'yyyy-mm-dd')")
    partitionrangestart = cur.fetchone()[0]
    cur.execute("select to_char(current_date+1,'yyyy-mm-dd')")
    partitionrangeend = cur.fetchone()[0]
    cur.execute("select count(*) from pg_tables where schemaname||'.'||tablename=lower('" + relation + "')")
    v_cnt = cur.fetchone()[0]
    if v_cnt:
        pass
    else:
        print 'ERROR: relation "' + relation + '" does not exist'
        return
    if reserve > 0:
        pass
    else:
        print 'ERROR: parameter "i_partition_add_cnt" must greater than 0,but you input ' + str(reserve)
        return

    cur.execute(
        "select pk_name_d,pk_name_d_suffix,p_d_parameter from gpmg.manager_table where schemaname||'.'||tablename=lower('" + relation + "')")
    ret = cur.fetchone()
    pk_name_d = ret[0]
    pk_name_d_suffix = ret[1]
    p_d_parameter = ret[2]
    sql_exists_partitions = "select 'alter table " + relation + " split default partition start ('||replace(replace(partitionrangestart,'timestamp without time zone','date'),' 00:00:00','')||') end ('||replace(replace(partitionrangeend,'timestamp without time zone','date'),' 00:00:00','')||') into (partition '||partitionname||',default partition);' from pg_partitions where partitionschemaname||'.'||tablename='" + relation + "' and to_date(substr(partitionrangestart,2,11),'yyyy-mm-dd')>='" + partitionrangestart + "'::date"
    cur.execute("select to_char(current_date,'" + pk_name_d_suffix + "')")
    p_name_d = pk_name_d + cur.fetchone()[0]
    sql_next_partitions = ""
    for i in range(1, reserve + 1):
        sql_next_partitions += "select 'alter table " + relation + " split default partition start (''" + partitionrangestart + "''::date) end (''" + partitionrangeend + "''::date) into (partition " + p_name_d + ",default partition);'"
        if i < reserve:
            sql_next_partitions += " union all "
        cur.execute("select to_char(current_date+" + str(i) + ",'" + pk_name_d_suffix + "')")
        pk_name_d_suffix_1 = cur.fetchone()[0]
        p_name_d = pk_name_d + pk_name_d_suffix_1
        cur.execute("select to_char(current_date+" + str(i) + ",'yyyy-mm-dd')")
        partitionrangestart = cur.fetchone()[0]
        cur.execute("select to_char(current_date+ " + str(i + 1) + ",'yyyy-mm-dd')")
        partitionrangeend = cur.fetchone()[0]
    cur.execute("select nextval('gpmg.seq_manager_part_mid')")
    manager_part_mid_sid = cur.fetchone()[0]
    cur.execute(
        "insert into gpmg.manager_part_mid(sid,t1) select " + str(
            manager_part_mid_sid) + ",* from (" + sql_exists_partitions + ") a ;")
    cur.execute(
        "insert into gpmg.manager_part_mid(sid,t2) select " + str(
            manager_part_mid_sid) + ",* from (" + sql_next_partitions + ") a ;")
    cur.execute(
        "insert into gpmg.manager_part_mid(sid,t3) select sid,t2 from gpmg.manager_part_mid where sid=" + str(
            manager_part_mid_sid) + " except select sid,t1 from gpmg.manager_part_mid where sid=" + str(
            manager_part_mid_sid) + ";")
    cur.execute("select t3 from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid) + " and t3 is not null")
    v_msg = "Done.Nothing."
    for record in cur.fetchall():
        sql_add_partitions = record[0]
        cur.execute(sql_add_partitions)
        v_msg = v_msg.replace('Nothing.', '')
    cur.execute("delete from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid))
    cur.close()
    db.commit()
    db.close()
    return v_msg


def partition_drop():
    db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
    cur = db.cursor()
    cur.execute("select  to_char(now(),'yyyymmddhh24miss');")
    log_id = cur.fetchone()[0]
    sql = "select schemaname,tablename,partlevel,p_type,pk_name_m,pk_name_d, pk_name_h, pk_name_m_suffix, pk_name_d_suffix, pk_name_h_suffix, retention_m, retention_d, retention_h from gpmg.manager_table where debug=0 and (status=1 or (status=0 and expire>current_date) ) and (oper_type='del' or oper_type='all') order by id"
    cur.execute(sql)
    serl = 0
    if len(sys.argv) == 2:
        ds = int(sys.argv[1])
    else:
        ds = 0
    v_delta = datetime.timedelta(days=ds)
    v_current_date = (datetime.date.today() + v_delta).strftime('%Y-%m-%d')

    for ret in cur.fetchall():
        v_schemaname = ret[0]
        v_tablename = ret[1]
        v_object = v_schemaname + "." + v_tablename
        sql_insert_log = "insert into gpmg.manager_table_operlog select '" + v_schemaname + "','" + v_tablename + "', now(),'d','" + log_id + "';"
        sql_update_log = "update gpmg.manager_table set drop_result=now() where schemaname||'.'||tablename='" + v_object + "';"
        serl += 1
        print "|" + str(datetime.datetime.now()).center(25) + "|" + str(serl).center(5) + "|" + v_object.center(
            45) + "|"

        v_partlevel = ret[2]
        v_p_type = ret[3]
        v_pk_name_m = ret[4]
        v_pk_name_d = ret[5]
        v_pk_name_h = ret[6]
        v_pk_name_m_suffix = ret[7]
        v_pk_name_d_suffix = ret[8]
        v_pk_name_h_suffix = ret[9]
        v_retention_m = ret[10]
        v_retention_d = ret[11]
        v_retention_h = ret[12]
        sql2 = "select count(*) from pg_tables where schemaname||'.'||tablename = '" + v_object + "'"
        cur.execute(sql2)
        is_exists = cur.fetchone()[0]
        if is_exists == 0:
            print "Not exists this table!\n"
            continue
        sql2 = "select count(1) from pg_locks where mode in ('AccessExclusiveLock','ExclusiveLock','RowExclusiveLock') and relation = '" + v_object + "'::regclass and gp_segment_id='-1' and granted=True;"
        is_lock = 0
        for i in range(5):
            cur.execute(sql2)
            is_lock = cur.fetchone()[0]
            if is_lock == 0:
                if i > 0:
                    print ''
                break
            print " sleep 10's ",
            time.sleep(10)
        if is_lock > 0:
            print "Can't lock this table!"
            continue

        if (',' + v_p_type + ',').find(',m,') >= 0:
            v_retention = str(v_retention_m)
            v_pk_name = v_pk_name_m
            v_pk_name_suffix = v_pk_name_m_suffix
            sql = "select '" + v_pk_name + "'||to_char('" + v_current_date + "'::date-interval '" + v_retention + " months' ,'" + v_pk_name_suffix + "')";
            cur.execute(sql)
            v_partitionname = cur.fetchone()[0]
            sql = "select count(*) from pg_partitions where schemaname||'.'||tablename='" + v_object + "' and partitionname='" + v_partitionname + "'";
            cur.execute(sql)
            is_exists = cur.fetchone()[0]
            if is_exists:
                sql = "select schemaname||'.'||partitiontablename from pg_partitions where schemaname||'.'||tablename='" + v_object + "' and partitionname='" + v_partitionname + "'";
                cur.execute(sql)
                partitiontablename = cur.fetchone()[0]
                cur.execute("select pg_total_relation_size('" + partitiontablename + "')")
                retsize = cur.fetchone()[0]
                cur.execute("select nextval('gpmg.seq_unload_size')")
                retid = cur.fetchone()[0]
                cur.execute(
                    "insert into gpmg.unload_size(id,rq,tablename,operation,tablesize) values(%s,%s,%s,'drop table',%s)",
                    (retid, datetime.date.today(), partitiontablename, retsize))
                cur.execute("alter table " + v_object + " drop partition " + v_partitionname)
                cur.execute(sql_update_log)
                cur.execute(sql_insert_log)
        if (',' + v_p_type + ',').find(',d,') >= 0:
            v_retention = str(v_retention_d)
            v_pk_name = v_pk_name_d
            v_pk_name_suffix = v_pk_name_d_suffix
            sql = "select '" + v_pk_name + "'||to_char('" + v_current_date + "'::date-" + v_retention + ",'" + v_pk_name_suffix + "')";
            cur.execute(sql)
            v_partitionname = cur.fetchone()[0]
            sql = "select count(*) from pg_partitions where schemaname||'.'||tablename='" + v_object + "' and partitionname='" + v_partitionname + "'";
            cur.execute(sql)
            is_exists = cur.fetchone()[0]
            if is_exists:
                sql = "select schemaname||'.'||partitiontablename from pg_partitions where schemaname||'.'||tablename='" + v_object + "' and partitionname='" + v_partitionname + "'";
                cur.execute(sql)
                partitiontablename = cur.fetchone()[0]
                cur.execute("select pg_total_relation_size('" + partitiontablename + "')")
                retsize = cur.fetchone()[0]
                cur.execute("select nextval('gpmg.seq_unload_size')")
                retid = cur.fetchone()[0]
                cur.execute(
                    "insert into gpmg.unload_size(id,rq,tablename,operation,tablesize) values(%s,%s,%s,'drop table',%s)",
                    (retid, datetime.date.today(), partitiontablename, retsize))
                cur.execute("alter table " + v_object + " drop partition " + v_partitionname)
                cur.execute(sql_update_log)
                cur.execute(sql_insert_log)
        db.commit()
    db.commit()
    db.close()


if __name__ == "__main__":
    pass
